import time
from copy import deepcopy
from get_sample import get_exp_sample2,get_template_sample
from Seacher.DocSearcher import DocSearcher
from utils.QA import QA
from prompt import system_prompt_orm
from llm import llm
from utils.pre_process_question import process_question
from kernel import CodeKernel, run_code
import re
import json


doc_searcher = DocSearcher('data/数据字典.xlsx')
orm_sample = get_exp_sample2()
template_sample = get_template_sample()

def init_kernel():
    ckernel = CodeKernel()
    ckernel.start()

    run_code('''
from IPython.core.interactiveshell import InteractiveShell
InteractiveShell.ast_node_interactivity = "last_expr"
# 简化异常显示
%xmode Minimal
''', True)
    run_code('''
import warnings
warnings.filterwarnings("ignore")
import pandas as pd
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', 20)
pd.set_option('display.float_format', lambda x: str(x))
from Models import *
from sqlalchemy import *
from utils.utils import exec_orm_sql
engine = create_engine('mysql+pymysql://')
''', ckernel, True)
    for _ in range(10):
        time.sleep(1)
        tr = run_code('''print("test output")''', ckernel, True)
        if tr == 'test output':
            break
    else:
        ckernel.shutdown()
        raise ValueError
    return ckernel


def FantasticORM2(question_dict):
    tid = question_dict['tid']
    try:
        qa = QA()
        ckernel = init_kernel()
        market_msg = ''
        msg = ''

        history = []
        for inx in range(len(question_dict['team'])):
            market_dic = {'ConstantDB.SecuMain': "A股", 'ConstantDB.HK_SecuMain': "港股",
                          'ConstantDB.US_SecuMain': "美股"}
            if inx == 0:
                info, tables = process_question(question_dict['team'][inx]['question'])
                if info:
                    msg = f'\n系统消息：已知的预处理信息为：`{info}`你可以使用以上信息继续查询。'
                if tables:
                    for table in tables:
                        market_msg += market_dic.get(table, '')

            question_info = qa.get_new_question(market_msg + question_dict['team'][inx]['question'])
            question = question_info['new_question']
            topic_type = question_info['type']

            qa.add_q(question)
            sub_tid = question_dict['team'][inx]['id']
            save_dict = {"raw_question":question_dict['team'][inx]['question'],
                        "question": question,'tid':tid,'sub_tid':sub_tid,
                         'final_code' :[]}

            exps = orm_sample.get_sample(question, 5)#.reverse()
            temp = template_sample.get_sample(question, 1)[0]

            dynamics_prompt = [
                {'role': 'system',
                 'content': system_prompt_orm.replace('[template]',temp)
                 }
            ]

            exps = '\n---\n'.join(exps)

            dynamics_prompt.extend(
                [{"role": 'user',
                  "content": f"---doc---{exps}\n---\n以上是一些ORM编写示例，你的代码参考此风格，但是对于新问题你要参考精细化示例重新思考，并严格按照`行为模式`执行。"},
                 {"role": "assistant", "content": "好的我已知晓以上是示例，我会再次独立思考。我以后会按照精细化示例的风格进行分析。并按照`行为模式`执行。"}])


            current_question = []
            if topic_type == '追加过滤条件':
                tip1 = '注意此问题是对之前结果(非示例)的追问，请注意这些结果的字段名称及其所在表，你可以直接使用或者修改刚才的查询变为子查询，或者使用之前查询增加外键字段 然后使用 `in_` 进行简化'
            elif topic_type == '结果追问':
                tip1 = '注意此问题是对之前结果(非示例)的追问，请注意这些结果的字段名称及其所在表'
            else:
                tip1 = ''

            if inx == 0 or topic_type == '继承首题实体的新问题':
                msg1 = msg

            else:
                msg1 = ''
            doc_txt = doc_searcher.search_by_question(question)
            doc_dict = {}
            # print(doc_txt)
            for d in doc_txt.split('\n---\n'):
                k = d.split('\n')[0].split(':')[1]
                doc_dict[k] = d
            current_question.extend([{'role': 'user', 'content': '请解决新问题：\n' + question + msg1 +'\n注意`思考`，注意`SQLAlchemy 编写准则`'},
                                     {'role': 'assistant', 'content': f'''<thought>基于文档查找手册，首先我要基于用户问题搜索文档</thought>
```python
doc_searcher.search_by_question("{question}")
```
请执行搜索并返回结果。
'''}, {"role": "user", "content": 'IPython:' + doc_txt +f'''
以上是文档，请分析文档以及用户的问题，选择正确的表格及其字段编写代码回答问题，注意表格和字段说明，如果问题比较复杂可以分多步解决，灵活使用orm和python。
{tip1}
步骤说明：
1. 举查询方法：
    - 列举两种可能的查询方法。优先方法：xxx 备用方法：xxx
2. 实现优先方法：
    - 编写代码实现优先方法。
    - 如果遇到困难则更换为备用方法。
3. 单表查询：
    - 主要关注单表查询。
    - 对于过滤条件，分析其必要性。
    - 对于不确定的过滤条件，选择相关过滤字段后，在pandas中进行进一步过滤。
4. 时间和筛选字段：
    - 特别注意时间和筛选字段的分析。
    - 对于历史数据（如股票数据），考虑交易日选择。
    - 对于其他查询，考虑最新发布的数据。
    - 公司状态如网址地址信息无需考虑。
5. 统计类信息：
    - 考虑筛选条件是否适用。
    - 如不确定，可以都统计，但在pandas中进行过滤。
6. 查询字段：
    - 查询xxxName字段时，如果有xxxCode字段，需同时查询xxxCode，例如查询SubclassName时顺带查询SubclassCode。
    - 思考需要查询的字段，并列举。
    - 对于日期过滤，默认为 EndDate，除非有其他合适的 Date
    - 对于股票计算，如涨跌幅，月年最高等，使用表内字段统计的，不要自己统计。 
7. 编码简化：
    - 尽量简化编码，大部分问题单表查询加上一些Python处理即可解决。
    - `exec_orm_sql`返回的是`DataFrame`，判断使用 `not result.empty`。
    - 所有models, sqlalchemy, pandas,engine 已经导入，无需重复import。请分析后编写代码，我将帮你执行代码。
请开始编码：
'''
       }])

            for try_num in range(7):
                if topic_type in ['继承首题实体的新问题','独立新问题']:
                    res = llm(dynamics_prompt + current_question)
                else:
                    res = llm(dynamics_prompt +history+ current_question)
                current_question.append({'role': 'assistant', 'content': res})
                if ('答案' in res  or '任务完成' in res ) and '```python' not in res:


                    question_dict['team'][inx]['answer_for_summary'] += res + current_question[-2]['content'].split('请判断输出结果能否回答问题')[0]
                    save_dict['messages'] = deepcopy(current_question)
                    qa.add_a(res)
                    save_dict['answer'] = res
                    #  ['search','search res','orm prompt','orm','answer prompt','answer']
                    cache_li = current_question[:2] + current_question[-4:]
                    cache_li[2]['content'] = '\n'.join([doc_dict[i] for i in doc_dict.keys() if i in '\n'.join(save_dict['final_code'])])
                    cache_li[2]['content'] = re.sub(r'"注释":.*"[,\b]]','', cache_li[2]['content'] )
                    cache_li[-2]['content'] = '...'+ cache_li[-2]['content'][-700:]

                    history.extend(cache_li)
                    try:
                        with open('/app/devlop_result/log.jsonl', 'a',encoding='utf8') as f:
                            f.write(json.dumps(save_dict,ensure_ascii=False))
                            f.write('\n')
                    except:
                        ...
                    break

                elif '【答案' in res and '```python' in res and 'doc_searcher.search_by_question' in res:
                    question_dict['team'][inx]['answer_for_summary'] += res
                code_list = re.findall(r'```python(.*?)```', res, re.DOTALL)
                run_log = []
                for code in code_list:
                    print('执行函数')
                    if 'doc_searcher.choose' in code:
                        code = re.search("doc_searcher.*$",code,re.DOTALL).group()
                        cache_output = []
                        save_dict['choose'] = code
                        try:
                            exec(f'r11 = {code}\n\ncache_output.append(r11)')
                        except Exception as e:
                            cache_output.append(str(e))
                        run_log.extend(cache_output)

                    elif 'doc_searcher' in code:
                        code = re.search('doc_searcher.*',code).group()
                        cache_output = []
                        try:
                            exec(f'r11 = {code}\n\ncache_output.append(r11)')
                        except  Exception as e:
                            cache_output.append(str(e))
                        run_log.extend(cache_output)
                    else:
                        for _ in range(2):
                            r = run_code(code, ckernel, iscode=True)
                            save_dict['final_code'].append(code)
                            if r:
                                break
                            else:
                                ckernel.shutdown()
                                ckernel=init_kernel()


                        if 'Empty DataFrame' in r:
                            r += f'''检测到查询为空，请判断空值是否合理，如果合理请回答问题:{question}，如果不合理请重新构造查询，
重新构造查询请使用以下方法：
1. 检查日期是否转换,字符串查询使用 like
2. 减少条件逐步筛选，并分步打印，只留下一个最稳妥的筛选条件，然后在pandas里进行下一步筛选
3. 查询少量数据样例，通过pandas进一步筛选
4. 更换表，更换Date等过滤条件。
5. 考虑此问题数量为0

注意`思考`和`SQLAlchemy 编写准则`'''
                        elif r == '':
                            r += '请使用print打印信息'
                        run_log.append(r)
                if run_log == []:
                    run_log = ['错误:代码运行失败或者超时，请优化ORM，尽量使用表外连表分步骤进行。']
                if try_num > 3:
                    current_question.pop(3)
                    current_question.pop(3)
                if try_num > 6:
                    current_question.append({'role': 'user', 'content': 'IPython:' + '\n'.join(run_log) +
f'\n请基于所有历史记录尽可能回答问题：`{question}`的内容,如果无法回答任何一点，请回答`【答案：任务失败】`，否则回答`【答案：可以回答的部分】`'})

                else:

                    current_question.append({'role': 'user', 'content': 'IPython:\n---日志---' + '\n'.join(run_log) + f"""
---
请判断输出结果能否回答问题`{question}`。
如果符合要求，请按照`任务完成\n【答案:xxx】`的格式重新整理。注意答案应从日志中提取的原始内容，不要增加或减少任何符号（如数字千分位）。
忽略格式化输出的错误，只要信息足够回答问题即可。注意答案格式以及用户需求的输出格式。
答案应该只有查询到的`值`的部分。对于公司名称，如无明确要求，应该回答`全称(简称)`；如果有要求，请仅回答指定字段。保留小数和日期格式也要注意。
其中日志里的公司名称格式为`全称(简称)`，请仔细判断输出的内容，不要只相信日志答案部分。
如果感觉文档有误，可以通过```python
doc_searcher.search_by_query('关键词1;关键词2')
```
根据关键词搜索文档,传入参数是一个str。
如果存在错误或不符合常识，请回答`任务继续`并重新修改代码。修复后的代码写在```python ```内。
"""})

    except:
        import traceback
        traceback.print_exc()

    return



